package com.loklok.code.confuse.util;

import com.intellij.openapi.application.ApplicationManager;
import com.intellij.openapi.progress.ProgressIndicator;
import com.intellij.openapi.progress.ProgressManager;
import com.intellij.openapi.progress.Task;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.ui.Messages;
import org.jetbrains.annotations.NotNull;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Consumer;
import java.util.function.Function;

public class BackgroundTaskRunner<Input, Result> {

    private final Project project;
    private final String title;
    private final List<Input> items;
    private final Function<Input, Result> backgroundWorker;
    private final Consumer<Result> uiConsumer;
    private final Runnable onCompletion;

    public BackgroundTaskRunner(Project project,
                                String title,
                                List<Input> items,
                                Function<Input, Result> backgroundWorker,
                                Consumer<Result> uiConsumer,
                                Runnable onCompletion) {
        this.project = project;
        this.title = title;
        this.items = items;
        this.backgroundWorker = backgroundWorker;
        this.uiConsumer = uiConsumer;
        this.onCompletion = onCompletion;
    }

    public void run() {
        ProgressManager.getInstance().run(new Task.Backgroundable(project, title, false) {
            private volatile ExecutorService executor;
            private final AtomicBoolean cancelled = new AtomicBoolean(false);
            private final List<Future<?>> futures = new ArrayList<>();

            @Override
            public void run(@NotNull ProgressIndicator indicator) {
                int total = items.size();
                CountDownLatch latch = new CountDownLatch(total);

                executor = Executors.newFixedThreadPool(Math.min(4, Runtime.getRuntime().availableProcessors()));

                indicator.setIndeterminate(false);
                indicator.setFraction(0.0);
                indicator.setText2(String.format("Processed %d of %d", 0, total));

                for (Input item : items) {
                    if (cancelled.get()) {
                        break;
                    }
                    Future<?> future = executor.submit(() -> {
                        try {
                            Result result = backgroundWorker.apply(item);
                            if (result != null && !cancelled.get()) {
                                ApplicationManager.getApplication().invokeLater(() -> {
                                    if (!cancelled.get()) {
                                        uiConsumer.accept(result);
                                    }
                                });
                            }
                        } catch (Exception e) {
                            if (!cancelled.get()) {
                                ApplicationManager.getApplication().invokeLater(() ->
                                        Messages.showErrorDialog(project, e.getMessage(), "Error"));
                            }
                        } finally {
                            latch.countDown();
                            updateProgress(indicator, latch.getCount(), total);
                        }
                    });

                    futures.add(future);
                }


                ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
                AtomicInteger processed = new AtomicInteger(0);

                // 定期更新进度条
                ScheduledFuture<?> progressUpdater = scheduler.scheduleAtFixedRate(() -> {
                    int current = processed.get();
                    double fraction = (double) current / total;
                    ApplicationManager.getApplication().invokeLater(() -> {
                        if (!cancelled.get()) {
                            indicator.setFraction(fraction);
                            indicator.setText2(String.format("Processed %d of %d", current, total));
                        }
                    });
                }, 0, 100, TimeUnit.MILLISECONDS); // 每100ms更新一次

                try {
                    while (!latch.await(100, TimeUnit.MILLISECONDS)) {
                        if (indicator.isCanceled()) {
                            cancelAllTasks();
                            break;
                        }
                    }
                } catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                }  finally {
                    progressUpdater.cancel(true);
                    scheduler.shutdownNow();
                }

                if (!cancelled.get() && onCompletion != null) {
                    ApplicationManager.getApplication().invokeLater(onCompletion);
                }

                executor.shutdown();
            }


            @Override
            public void onCancel() {
                cancelAllTasks();
            }

            private void cancelAllTasks() {
                cancelled.set(true);
                if (executor != null) {
                    executor.shutdownNow();
                }
                futures.forEach(f -> f.cancel(true));
            }

            private void updateProgress(ProgressIndicator indicator, long remaining, int total) {
                if (cancelled.get()) {
                    return;
                }

                int progress = total - (int) remaining;
                double fraction = (double) progress / total;

                ApplicationManager.getApplication().invokeLater(() -> {
                    if (!cancelled.get()) {
                        indicator.setFraction(fraction);
                        indicator.setText2(String.format("Processed %d of %d", progress, total));
                    }
                });
            }
        });
    }


}
